/************************************************************************
Copyright 2003, 2004, 2005, 2006, 2007, 2008 Grid Technology Research Center,
National Institute of Advanced Industrial Science and Technology.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
************************************************************************/

package jp.aist.gtrc.plus.scheduler;

import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;

import org.apache.log4j.Logger;

import jp.aist.gtrc.plus.scheduler.node.NodeAllocator;
import jp.aist.gtrc.plus.scheduler.node.NodeInfo;
import jp.aist.gtrc.plus.scheduler.node.NodeSorter;
import jp.aist.gtrc.plus.scheduler.specific.sge.jaxb.common.JGElement;
import jp.aist.gtrc.plus.scheduler.specific.sge.sched.SGESelfExecNodeAllocator;
import jp.aist.gtrc.plus.scheduler.specific.sge.sched.SGESelfMainServer;
import jp.aist.gtrc.plus.scheduler.specific.sge.sched.SGESelfReserveNodeAllocator;
import jp.aist.gtrc.plus.scheduler.specific.sge.sched.SGESelfStatusManager;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEJobStatus;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEQInstance;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGETaskStatus;
import jp.aist.gtrc.plus.scheduler.status.JobStatus;
import jp.aist.gtrc.plus.scheduler.status.MainServer;
import jp.aist.gtrc.plus.scheduler.status.StatusManager;
import jp.aist.gtrc.plus.scheduler.sys.SchedulerException;

class SGESelfScheduler extends DefaultSelfScheduler {

        protected static final Logger logger = Logger.getLogger(SGESelfScheduler.class);

	private boolean jobReruned;

	SGESelfScheduler(SchedulerOptions options) throws Exception {

		super(options);
		startReserveServer();

	}

	protected StatusManager getStatusManager(SchedulerOptions options) throws SchedulerException {

		return SGESelfStatusManager.getInstance();

	}

	protected NodeAllocator getExecNodeAllocator(SchedulerOptions options) throws SchedulerException {

		return new SGESelfExecNodeAllocator(new NodeSorter(options));

	}

	protected NodeAllocator getRsvNodeAllocator(SchedulerOptions options) throws SchedulerException {

		return new SGESelfReserveNodeAllocator(options, nodeMgr, statusMgr);

	}

	protected MainServer getMainServer() throws Exception  {

		return (MainServer)SGESelfMainServer.getInstance();

	}

	protected void runSchedule(ScheduleOrder order) throws SchedulerException {

		super.runSchedule(order);
		afterSchedule(order.getMainServer());

	}

	protected void schedule(ScheduleOrder order) throws SchedulerException{
                logger.debug("SGESelfScheduler schedule()");

		MainServer srv = order.getMainServer();

		Collection<JobStatus> allJobs = statusMgr.getJobs();
		Collection<JobStatus> runnableJobs = jobMgr.getRunnableJobs(allJobs);
		if (runnableJobs == null)
			return;

		jobReruned = false;
		while ((runnableJobs.size() > 0) && (jobReruned == false)) {
			JobStatus job = scheduler.selectJob(runnableJobs);
			runOneJob(srv, job);
			runnableJobs.remove(job);
		}

	}

	protected boolean handleUnreservedJobs(MainServer srv, Collection<NodeInfo> rsvdNodes) {

		Collection<JobStatus> allTasks = statusMgr.getJobs();
		HashSet<SGEJobStatus> reranJobs = new HashSet<SGEJobStatus>();

		/*
		 * 'tasksOnNodeMap' is HashMap of Task list
		 * ex.
		 * 	 key		value [ task list of running on the node ]
		 * node00		{ TaskA, TaskB }
		 * node01		{ TaskA, TaskX, taskY }
		 */
		HashMap<String, LinkedList<SGETaskStatus>> tasksOnNodeMap
			= new HashMap<String, LinkedList<SGETaskStatus>>();
		for (JobStatus t : allTasks) {
			SGETaskStatus task = (SGETaskStatus)t;
			if (task.isRunning() == false) {
				continue;
			}
			if (task.getReserveId() != null) {
				// 'task' is running with reserve. OK to run.
				continue;
			}
			for (JGElement gq : task.getGrantedQueues()) {
				String nodename = gq.getJGQhostname();
				LinkedList<SGETaskStatus> list = tasksOnNodeMap.get(nodename);
				if (list == null) {
					list = new LinkedList<SGETaskStatus>();
					tasksOnNodeMap.put(nodename, list);
				}
				list.add(task);
			}
		}
		if (tasksOnNodeMap.size() == 0)
			return false;

		for (NodeInfo node : rsvdNodes) {
			SGEQInstance qIns = (SGEQInstance)node.getStatus();

			// NOTE: 'nodename' is queue fullname actually
			String nodename = qIns.getName();
			Collection<SGETaskStatus> tasksOnNode = tasksOnNodeMap.get(nodename);
			if (tasksOnNode == null) {
				// no unreserved-running task on this 'node'
				continue;
			}
			for (SGETaskStatus task : tasksOnNode) {
				SGEJobStatus job = task.getSGEJobStatus();
				if (reranJobs.contains(job)) {
					// already reran 'job' by previous task
					continue;
				}

				try {
					/*
					 * NOTE: rerunJob(SGETaskStatus) will
					 * 	 1. resubmit 'job' (not 'task'), make clone of 'job'
					 * 	 2. delete (means kill) 'job'
					 * After them, clone of 'job' remains.
					 */
					srv.rerunJob(task);
					jobReruned = true;
				} catch (SchedulerException e) {
					logger.info("cannot rerun of " + task.getName());
				}

				// 'job' is already reran, not check anymore
				reranJobs.add(job);
			}
		}

		return jobReruned;

	}

	protected void afterSchedule(MainServer srv)  {

		SGESelfMainServer sgeSrv = (SGESelfMainServer)srv;
		Collection<JobStatus> allJobs = statusMgr.getJobs();
		for (JobStatus j : allJobs) {
			SGETaskStatus task = (SGETaskStatus)j;
			if (task.isReadyToDelete()) {
				try {
					sgeSrv.expireJob(task.getJobID());
				} catch (SchedulerException e) {
					// ignore fail
					logger.info("Failed to expire job: " + task.getJobID());
				}
			}
		}

	}

	protected void stopScheduler() {

		try {
			SGESelfMainServer srv = SGESelfMainServer.getInstance();
			srv.shutdown();
		} catch (SchedulerException e) {
			// ignore
		}

		super.stopScheduler();

	}

}
